#!/usr/bin/python

"""
This program draws gene clusters to make publication quality figures. It takes in
Genbank file, COG category file, and expects start and stop coordinates of region
to inspect.

Usage: dissertation_DrawGeneswithPtt.py seq.gbk cogs.t.list 10000 20000
Examples:
dissertation_DrawGeneswithPtt.py ../../../annotation/Genome.gbf Genome.ptt cogs.t.list 10000 20000

Note: Resolution is best if the segment in view is less than 50000bp.

Author: Jimmy Saw
Date of last update: 04-23-2012

"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC
import matplotlib.patches as mpatch

#Regex and other stuffs
cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')
pttcog = re.compile('(COG\d{4})(\w).*')

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#8B7500',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
    '-' : '#D3D3D3'
    }

##Genome one Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    if feat.type == 'CDS':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 1 ptt file
g1cogs = {}
g1pttfile = open(sys.argv[2], "rU")
g1ptt = g1pttfile.readlines()
for line in g1ptt[3:]:
    c = line.split('\t')
    g1ltag = c[5]
    g1gene = c[4]
    #g1cogcat = '-'
    g1cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        #g1cogcat = p.group(2)
        g1cog = p.group(1)
    #g1cogs[g1ltag] = g1cogcat
    g1cogs[g1ltag] = g1cog

cogcatfile = open(sys.argv[3], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

cogcatfile.close()

spanx1 = int(sys.argv[4])
spanx2 = int(sys.argv[5])

glist = []

genome1x = [spanx1, spanx2]
genome1y = [2, 2]
genome2x = [spanx1, spanx2]
genome2y = [10, 10]
mid1x = [spanx1+200, spanx2-200]
mid1y = [6.75, 6.75]

##Start plotting
fig = plt.figure(1, figsize=(16,4))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
#ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
#    mec='black', ls='-', lw=2.0)
#ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
#    mec='black', ls='-', lw=2.0)
ax1.plot(mid1x, mid1y, color='#CDAA7D', marker='|', mec='#CDAA7D', ls =':', lw=2.0)

#ax1.axis([0, g1length, 0, 14])
ax1.axis([spanx1, spanx2, 0, 14])

for k, v in g1featdict.iteritems():
    g1feat = v
    g1locustag = g1feat.qualifiers['locus_tag'][0]
    g1start = g1feat.location._start.position
    g1stop = g1feat.location._end.position
    g1size = g1stop - g1start + 1
    g1mid = g1start + ((g1stop - g1start) / 2.0)
    g1desc = g1feat.qualifiers['product'][0]
    g1gene = ""
    if g1feat.qualifiers.has_key('gene'):
        g1gene = g1feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g1gene = g1feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = "#D3D3D3" #base color

    if g1locustag in g1cogs:
        if g1cogs[g1locustag] != '-':
            #cogcolor = cogdict[g1cogs[g1locustag]]
            cogcolor = cogdict[cogcatdict[g1cogs[g1locustag]]]

    if g1feat.type == 'tRNA':
        cogcolor = '#800000'
    if g1feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g1feat.strand == -1:
        if spanx1 <= g1start <= spanx2 and spanx1 <= g1stop <= spanx2:
            rect = Rectangle((g1start, 6.0), g1size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
            plt.gca().add_patch(rect)
            ax1.text(g1mid, 6.5, g1gene, fontsize=8, color='black', rotation=45)
            #ax1.plot(g1start, 6.0, 'r<', mec='red')
    else:
        if spanx1 <= g1start <= spanx2 and spanx1 <= g1stop <= spanx2:
            rect = Rectangle((g1start, 7.0), g1size, 0.5, fc=cogcolor, ec=cogcolor, alpha=0.5)
            plt.gca().add_patch(rect)
            ax1.text(g1mid, 7.5, g1gene, fontsize=8, color='black', rotation=45)
            #ax1.plot(g1stop, 7.0, 'r>', mec='red')

#Draw legend box for COG categories
coglist = []
for k, v in cogdict.iteritems():
    coglist.append((k,v))

coglist.sort()

ccounts = len(coglist)
a = 0
spansize = spanx2 - spanx1
spanmid = spanx1 + ((spanx2 - spanx1) / 2.0)
xstart = spanmid
#increment = 20000 #for synecoccus
increment = spansize * 0.01 #for gloeobacter
while a < ccounts:
    fc = coglist[a][1]
    tx = coglist[a][0]
    #rect = Rectangle((xstart, 2.5), 20000, 0.25, facecolor=fc, alpha=0.5) #for synecococcus
    rect = Rectangle((xstart, 2.5), increment, 0.5, fc=fc, alpha=0.5) #for gloeobacter
    plt.gca().add_patch(rect)
    ax1.annotate(tx, xy=(xstart+(increment/2.0), 2.2), horizontalalignment='center', verticalalignment='center', fontsize=8)
    a += 1
    xstart = xstart + increment

ax1.annotate('COG categories', xy=(spanmid, 1.2), horizontalalignment='left', verticalalignment='center', fontsize=10)
ax1.annotate(g1seq.annotations['organism'], xy=(0.5, 0.9), xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)

"""
#This segment is an example of how to draw polygon connecting genomic regions
a = [140000, 140000, 141000-(1000 * 0.1), 141000] # + strand
b = [5, 5.5, 5.5, 5]
c = [141000+(1000*0.1), 141000, 142000, 142000] # - strand
d = [4.5, 5, 5, 4.5]

ax1.fill(a, b, fc='red', ec='red', alpha=0.2)
ax1.fill(c, d, fc='blue', ec='blue', alpha=0.2)
"""


"""
ax1.text(1710000, 3, 'Test',
          ha="center",
          size=30,
          bbox=dict(boxstyle='rarrow', fc="w", ec="k"))
"""


frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)


plt.show()


